import cv2
import face_recognition
import pickle
import numpy as np
import os
import base64
from flask import Flask, render_template, request, jsonify, redirect, url_for

app = Flask(__name__)

# 确保face_data目录存在
if not os.path.exists('face_data'):
    os.makedirs('face_data')

@app.route('/')
def index():
    """主页"""
    return render_template('index.html')

@app.route('/train')
def train_page():
    """人脸训练页面"""
    return render_template('train.html')

@app.route('/unlock')
def unlock_page():
    """人脸识别解锁页面"""
    if not os.path.exists('face_data/user_face.pkl'):
        return redirect(url_for('no_data_page'))
    return render_template('unlock.html')

@app.route('/no_data')
def no_data_page():
    """无数据页面"""
    return render_template('no_data.html')

@app.route('/api/train', methods=['POST'])
def train_face():
    """API接口：训练人脸识别"""
    data = request.get_json()
    image_data = data['image']
    
    # 解码图像数据
    header, encoded = image_data.split(',', 1)
    image_bytes = base64.b64decode(encoded)
    
    # 将图像数据转换为numpy数组
    nparr = np.frombuffer(image_bytes, np.uint8)
    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    
    # 检测人脸
    face_locations = face_recognition.face_locations(img)
    
    if len(face_locations) > 0:
        encoding = face_recognition.face_encodings(img, face_locations)[0]
        
        # 读取已有的人脸编码或创建新列表
        if os.path.exists('face_data/user_face.pkl'):
            with open('face_data/user_face.pkl', 'rb') as f:
                face_encodings = pickle.load(f)
        else:
            face_encodings = []
        
        face_encodings.append(encoding)
        
        # 保存人脸编码
        with open('face_data/user_face.pkl', 'wb') as f:
            pickle.dump(face_encodings, f)
        
        return jsonify({
            'success': True,
            'message': f'成功采集样本，当前共有 {len(face_encodings)} 个样本',
            'sample_count': len(face_encodings)
        })
    else:
        return jsonify({
            'success': False,
            'message': '未检测到人脸，请确保面部在摄像头范围内'
        })

@app.route('/api/get_samples_count')
def get_samples_count():
    """API接口：获取已采集样本数量"""
    if os.path.exists('face_data/user_face.pkl'):
        with open('face_data/user_face.pkl', 'rb') as f:
            face_encodings = pickle.load(f)
        return jsonify({
            'success': True,
            'count': len(face_encodings)
        })
    else:
        return jsonify({
            'success': True,
            'count': 0
        })

@app.route('/api/unlock', methods=['POST'])
def face_unlock():
    """API接口：人脸识别解锁"""
    # 检查数据文件
    if not os.path.exists('face_data/user_face.pkl'):
        return jsonify({
            'success': False,
            'message': '未找到人脸数据，请先进行人脸录入'
        })
    
    # 加载训练数据
    with open('face_data/user_face.pkl', 'rb') as f:
        known_encodings = pickle.load(f)
    
    data = request.get_json()
    image_data = data['image']
    
    # 解码图像数据
    header, encoded = image_data.split(',', 1)
    image_bytes = base64.b64decode(encoded)
    
    # 将图像数据转换为numpy数组
    nparr = np.frombuffer(image_bytes, np.uint8)
    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    
    # 检测人脸
    face_locations = face_recognition.face_locations(img)
    face_encodings = face_recognition.face_encodings(img, face_locations)
    
    if len(face_encodings) > 0:
        # 计算相似度
        distances = face_recognition.face_distance(known_encodings, face_encodings[0])
        min_distance = np.min(distances)
        threshold = 0.4  # 识别阈值
        
        # 判断是否匹配
        if min_distance < threshold:
            return jsonify({
                'success': True,
                'match': True,
                'distance': float(min_distance),
                'threshold': threshold,
                'message': '解锁成功！'
            })
        else:
            return jsonify({
                'success': True,
                'match': False,
                'distance': float(min_distance),
                'threshold': threshold,
                'message': '访问被拒绝'
            })
    else:
        return jsonify({
            'success': True,
            'face_detected': False,
            'message': '未检测到人脸'
        })

if __name__ == '__main__':
    app.run(debug=True, port=5001)